# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

import os
import pickle
import pandas
import numpy

from absl import flags

import coin.strategy.mm.tool.archive_base as abase
from coin.strategy.mm.tool.signal_plot_base import dump_book

from coin.strategy.mm.tool.order_reader import OrderReader


class OrderStrategy(abase.SerializedJob):
  def __init__(self, baseflags, osflags, product):
    self.baseflags = baseflags
    self.osflags = osflags
    self.trading_date = self.baseflags.trading_date
    self._product = product
    self._order_reader = OrderReader(self._product)

  def on_log(self, *args, **kwargs):
    self._order_reader.on_log2(*args, **kwargs)

  def set_time_range(self, *args, **kwargs):
    self._order_reader.set_time_range(*args, **kwargs)

  def get_filename_prefix(self):
    symbol = self._product.symbol if self._product is not None else ""
    return "_".join((self.trading_date,
                     self._product.exchange,
                     symbol,
                     self.osflags.orderlog_machine,
                     self.osflags.owner))


class OrderDumpJob(abase.SerializedJob):
  def __init__(self, job_dir, job_cls, baseflags, fsflags, osflags, product):
    self.job_cls = job_cls
    self.baseflags = baseflags
    self.fsflags = fsflags
    self.osflags = osflags
    self.product = product
    self.symbol = self.product.symbol if self.product is not None else ""

    if flags.FLAGS.use_pickle:
      os.makedirs(job_dir, exist_ok=True)
      os.makedirs(job_dir + "_pkl", exist_ok=True)
      abase.SerializedJob.__init__(
          self,
          os.path.join(
              job_dir + "_pkl",
              "%s.pkl"
              % "_".join([baseflags.trading_date, self.osflags.orderlog_machine, self.symbol])))

  def run(self):
    if flags.FLAGS.use_pickle:
      if not os.path.exists(self.filename) or flags.FLAGS.force_rerun:
        self.serialize(self.filename)
      strat, bbodf = self.deserialize()
      return strat, bbodf
    else:
      return self.serialize(None)

  def serialize(self, filename):
    tds = abase.get_trading_dates(self.baseflags.trading_date)
    assert len(tds) == 1
    trading_date = tds[0]
    self.baseflags.trading_date = trading_date.strftime("%Y%m%d")

    print(self.baseflags, self.osflags)
    start_time, end_time = abase.get_time_range(self.baseflags)

    strat = self.job_cls(self.baseflags, self.osflags, self.product)
    strat.set_time_range(start_time, end_time)

    if hasattr(strat, "on_strat_log"):
      abase.run_from_strat_archive_base(self.baseflags, self.osflags, on_log=strat.on_strat_log)
    else:
      abase.run_from_order_archive_base(self.baseflags, self.osflags, on_log=strat.on_log)

    bbos = dump_book(self.product, trading_date, flags.FLAGS.time_range)
    bbodf = pandas.DataFrame(bbos)
    sel_nan_trade = numpy.isnan(bbodf['trade_price'])
    for trade_col in bbodf.columns:
      if trade_col.find('trade') < 0:
        continue
      bbodf[trade_col] = bbodf[trade_col].fillna(method='ffill')
    bbodf = bbodf.loc[sel_nan_trade].reset_index(drop=True)
    idx_prc_same = (bbodf['bid0_price'].diff() == 0) & (bbodf['ask0_price'].diff() == 0)
    bbodf = bbodf.loc[~idx_prc_same].reset_index(drop=True)

    if flags.FLAGS.use_pickle:
      with open(filename, 'wb') as writefile:
        pickle.dump((strat, bbodf), writefile)
    else:
      return (strat, bbodf)

  def deserialize_impl(self, filename):
    return pickle.load(open(filename, 'rb'))
